from typing import Union, Callable, List, Optional

import paddle
import paddle.fluid.layers as L
from paddle.fluid.dygraph import Layer
from paddle.fluid.optimizer import Optimizer

from helm.dynamic.engine.engine import Engine
from helm.dynamic.engine.callback import DefaultTrainLogger, DefaultEvalLogger, ModelCheckpoint
from helm.dynamic.engine.metrics import Metric


def create_trainer(
        model: Layer,
        criterion: Union[Callable, Layer],
        optimizer: Optimizer,
        metrics: List[Metric],
        log_freq: int = 10,
        save_freq: Optional[int] = None,
        save_dir: Optional[int] = None,
):
    def train_batch(engine, batch):
        engine.model.train()

        lr, hr = batch

        pred = engine.model(lr)
        loss = engine.criterion(pred, hr)
        loss.backward()
        engine.optimizer.minimize(loss)
        engine.model.clear_gradients()

        output = {
            "loss": float(loss.numpy()),
            "y_true": hr,
            "y_pred": pred.detach(),
            "batch_size": lr.shape[0],
            "learning_rate": engine.optimizer.current_step_lr(),
        }

        return output

    callbacks = [*metrics, DefaultTrainLogger(log_freq)]
    if save_freq:
        callbacks.append(ModelCheckpoint(save_freq, save_dir))
    engine = Engine(
        train_batch, callbacks,
        model=model, criterion=criterion, optimizer=optimizer)
    return engine


def quantize(x):
    return L.round(L.clamp(x, 0, 255))


def create_evaluator(
        model: Layer,
        criterion: Union[Callable, Layer],
        metrics: List[Metric],
        log_freq: int = 10,
        stage: str = "Valid",
):

    def test_batch(engine, batch):
        engine.model.eval()

        lr, hr = batch

        with paddle.fluid.dygraph.no_grad():
            pred = engine.model(lr)
            loss = engine.criterion(pred, hr)

        output = {
            "loss": float(loss.numpy()),
            "y_true": hr,
            "y_pred": quantize(pred),
            "batch_size": lr.shape[0],
        }

        return output

    callbacks = [*metrics, DefaultEvalLogger(log_freq)]
    engine = Engine(test_batch, callbacks,
                    model=model, criterion=criterion, stage=stage)

    return engine